import os
import sys
import numpy as np
import pydicom
import matplotlib.pyplot as plt
import cv2

validRow = 100

# 计算标准化互相关系数（NCC）
def NCC(im1, im2, eps=1e-8):
    var1 = im1 - im1.mean()
    var2 = im2 - im2.mean()
    ncc = np.sum(var1 * var2) / (np.sqrt(np.sum(var1 ** 2)) * np.sqrt(np.sum(var2 ** 2)) + eps)
    return ncc

# 使用NCC进行图像配准
def NCCregistration(img1_r, img2_c):
    # startRow迭代初始位置
    startRow = img2_c.shape[0] // 2
    endRow = img2_c.shape[0] - validRow + 1
    best = {'startRow': startRow, 'cost': np.inf}

    for i in range(startRow, endRow):
        j = i + validRow
        moving_valid = img2_c[i: j, :]
        ncc = NCC(img1_r, moving_valid)
        cost = 1 - abs(ncc)
        if cost < best['cost']:
            best['startRow'] = i
            best['cost'] = cost
    return best['startRow']

def see(data):
    plt.imshow(data, cmap='gray')
    plt.show()

def read_dcms(path):
    files_names = os.listdir(path)
    files_names.sort(key=lambda D_name: int(D_name.split('.')[0].split('_')[-1]))
    # print(files_names, len(files_names))
    files_num = 0
    imgs = []
    for name in files_names:
        file_path = os.path.join(path, name)
        dic_tag = pydicom.read_file(file_path)
        img = dic_tag.pixel_array
        files_num = files_num + 1
        imgs.append(img)
    return imgs, files_num

def fusion(dcms, lens):
    # 第一张数据作为fixed数据并初始化
    fixed_img = dcms[0]

    # 循环读取moving数据
    for i in range(1, lens):
        moving_img = dcms[i]
        #计算固定图像的前validRow行与当前移动图像之间的最佳配准位置（best）
        best = NCCregistration(fixed_img[0: validRow, :], moving_img)
        #生成一个线性空间，用于在配准过渡区域上创建一个逐渐变化的权重（ramp）
        ramp = np.linspace(0, 1, validRow)
        #创建一个缓冲区（result_buffer），用于存储拼接后的图像数据。它的大小为固定图像高度加上移动图像配准位置（best）
        result_buffer = np.empty((fixed_img.shape[0] + best, fixed_img.shape[1]), dtype='uint16')
        #将移动图像的前best行复制到缓冲区的相应位置
        result_buffer[0: best, :] = moving_img[0: best, :]
        #对配准过渡区域进行加权平均，使用线性空间ramp对移动图像和固定图像的前validRow行进行混合
        result_buffer[best: best + validRow, :] = np.transpose(np.transpose(moving_img[best: best + validRow, :]) * (1 - ramp) +
                                                               np.transpose(fixed_img[: validRow, :]) * ramp)
        #将固定图像的剩余部分复制到缓冲区的相应位置
        result_buffer[best + validRow:, :] = fixed_img[validRow:, :]
        #将拼接后的图像数据赋值给fixed_img，为下一次循环提供固定图像。
        fixed_img = result_buffer
    return fixed_img

if __name__ == '__main__':
    OriData_path = './testData'
    print("********按序读入待拼接dcm数据*******")
    dcms, dcms_num = read_dcms(OriData_path)
    print("*********开始拼接数据*********")
    result = fusion(dcms, dcms_num)
    print('打印最终拼接数据的尺寸:', result.shape)
    see(result)
    cv2.imwrite("./image.jpg", result)
